import json
from django.core.exceptions import ValidationError
from django.db import models
from django.forms import model_to_dict
from django.urls import reverse
from django.contrib.auth.models import User
from django.contrib.staticfiles.storage import staticfiles_storage
from django.db.models import Q
from django.db.models.constraints import UniqueConstraint
from .utils import get_key_choices
from zhon.hanzi import punctuation
import string
import logging

# import spacy
# from spacy.tokens import Doc


# class WhitespaceTokenizer(object):
#     def __init__(self, vocab):
#         self.vocab = vocab
#
#     def __call__(self, text):
#         words = text.split(' ')
#         # All tokens 'own' a subsequent space character in this tokenizer
#         spaces = [True] * len(words)
#         return Doc(self.vocab, words=words, spaces=spaces)
logger = logging.getLogger(__name__)


class Project(models.Model):
    SEQUENCE_LABELING = 'SequenceLabeling'

    name = models.CharField(max_length=100, verbose_name="项目名称")
    description = models.CharField(max_length=500, verbose_name="项目描述")
    guideline = models.TextField()
    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)
    users = models.ManyToManyField(User, related_name='projects', verbose_name="用户")

    def get_absolute_url(self):
        return reverse('upload', args=[self.id])

    def get_progress(self):
        docs = self.get_unannotated_documents()
        total = self.documents.count()
        remaining = docs.count()
        return {'total': total, 'remaining': remaining}

    @property
    def image(self):
        url = staticfiles_storage.url('images/cat-3449999_640.jpg')
        return url

    def get_template_name(self):
        template_name = 'annotation/sequence_labeling.html'
        return template_name

    def get_documents(self):
        docs = self.documents.all()
        return docs

    def get_unannotated_documents(self):
        docs = self.documents.all()
        docs = docs.filter(annotated=False)
        return docs

    def get_index_documents(self, indices):
        docs = self.documents.all()
        #docs_indices = [d.id for d in docs]
        #active_indices = [docs_indices[i - 1] for i in indices]
        docs = list(docs.filter(pk__in=indices))
        docs.sort(key=lambda t: indices.index(t.pk))
        return docs

    def get_document_serializer(self):
        from .serializers import SequenceDocumentSerializer
        return SequenceDocumentSerializer

    def get_annotation_serializer(self):
        from .serializers import SequenceAnnotationSerializer
        return SequenceAnnotationSerializer

    def get_relation_annotation_serializer(self):
        from .serializers import RelationAnnotationSerializer
        return RelationAnnotationSerializer

    def get_attribute_annotation_serializer(self):
        from .serializers import AttributeAnnotationSerializer
        return AttributeAnnotationSerializer

    def get_annotation_class(self):
        return SequenceAnnotation

    def get_label_list(self,task):
        if task == 'ner':
            labels = self.labels.all()
            label_list = []
            for label in labels:
                label_list.append('B-' + str(label.text))
                label_list.append('I-' + str(label.text))
            label_list.append('O')
            label_list.append('X')
        else:
            labels = self.labels.all()
            relaion_labels = self.relation_labels.all()
            entities = []
            relations = []
            for label in labels:
                entities.append(label.text)
            for relation_label in relaion_labels:
                relations.append(relation_label.relation_name)
            label_list = {"entities": entities, relations: relations}
        return label_list

    def __str__(self):
        return self.name


class Label(models.Model):
    KEY_CHOICES = get_key_choices()
    COLOR_CHOICES = ()

    text = models.CharField(max_length=100)
    shortcut = models.CharField(max_length=15, blank=True, null=True, choices=KEY_CHOICES)
    project = models.ForeignKey(Project, related_name='labels', on_delete=models.CASCADE)
    background_color = models.CharField(max_length=7, default='#209cee')
    text_color = models.CharField(max_length=7, default='#ffffff')

    def __str__(self):
        return self.text

    class Meta:
        unique_together = (
            ('project', 'text'),
            ('project', 'shortcut')
        )


class Document(models.Model):
    text = models.TextField()
    project = models.ForeignKey(Project, related_name='documents', on_delete=models.CASCADE)
    annotated = models.BooleanField(default=False)
    related = models.BooleanField(default=False)
    submitted = models.BooleanField(default=False)
    annotating = models.BooleanField(default=False)
    metadata = models.TextField(default='{}')

    def delete_annotations(self):
        self.seq_annotations.all().delete()

    def delete_relation_annotations(self):
        self.seq_relation_annotations.all().delete()

    def delete_attribute_annotations(self):
        self.seq_attribute_annotations.all().delete()

    def get_annotations(self):
        seq_annotation = self.seq_annotations.all()
        return seq_annotation

    def get_relation_annotations(self):
        seq_relation_annotation = self.seq_relation_annotations.all()
        return seq_relation_annotation

    def get_attribute_annotations(self):
        seq_attribute_annotation = self.seq_attribute_annotations.all()
        return seq_attribute_annotation

    def get_annotations_user(self, user_id):
        seq_annotation = self.seq_annotations.filter(document_id=self.id, user_id=user_id)
        return seq_annotation

    def to_csv(self, user_id):
        return self.make_dataset(user_id)

    def make_dataset(self, user_id):
        return self.make_dataset_for_sequence_labeling(user_id)

    def make_relation_data(self):
        # TODO:如何处理Other这种情况?
        # TODO:是否存在删除关系没删完
        res = []
        relation_annotations = RelationAnnotation.objects.filter(document_id=self.id)
        text = self.text.strip()
        text = [d for d in self.text]
        for relation in relation_annotations:
            from_label = relation.from_label
            to_label = relation.to_label
            relation_label = relation.relation_label
            direction = 0
            if from_label.start_offset > to_label.start_offset:
                direction = 1
            if direction == 0:
                tmp = ((from_label.start_offset, from_label.end_offset),
                       (to_label.start_offset, to_label.end_offset))
            else:
                tmp = ((to_label.start_offset, to_label.end_offset),
                       (from_label.start_offset, from_label.end_offset),)
            one = text[:tmp[0][0]] + ["<e1>"] + text[tmp[0][0]:tmp[0][1]] + ["</e1>"] \
                  + text[tmp[0][1]:tmp[1][0]] + ["<e2>"] + text[tmp[1][0]:tmp[1][1]] + ["</e2>"] \
                  + text[tmp[1][1]:]
            item = {"sentence": one, "relation": relation_label.relation_name + "$" + str(direction)}
            print(item, direction)
            res.append(item)
        return res

    def make_dataset_for_sequence_labeling(self, user_id):
        annotations = self.get_annotations_user(user_id=user_id)
        words = [token for token in self.text]
        dataset = [[self.id, word, 'O', self.metadata] for word in words]

        def is_valid(item):
            word = item[1]
            if word == " ":
                return False
            if word in punctuation or word in string.punctuation:
                return False
            return True

        for a in annotations:
            if a.label:
                dataset[a.start_offset][2] = 'B-{}'.format(a.label.text)
                for i in range(a.start_offset + 1, a.end_offset):
                    dataset[i][2] = 'I-{}'.format(a.label.text)
        # dataset = [item for item in dataset if is_valid(item)]
        return dataset


    def make_train_data(self, task):
        ner_annotations = self.get_annotations()
        re_annotations = self.get_relation_annotations()
        assert task in ['ner','re','ner_re']
        if task=='ner':
            words = [token for token in self.text]
            text = ' '.join(words)
            labels = ['O' for word in words]

            for a in ner_annotations:
                if a.label:
                    labels[a.start_offset] = 'B-{}'.format(a.label.text)
                    for i in range(a.start_offset + 1, a.end_offset):
                        labels[i] = 'I-{}'.format(a.label.text)
            data = {'text': text, 'labels': labels}
        else:
            entities=[]
            relations=[]
            words = [token for token in self.text]
            #
            entity_mapping={}
            entity_count=0
            for a in ner_annotations:
                if a.label:
                    entity={'type':a.label.text,'start':a.start_offset,'end':a.end_offset}
                    entities.append(entity)
                    entity_mapping[a.id]=entity_count
                    entity_count+=1
            for a in re_annotations:
                relation={'type':a.relation_label.relation_name,'head':entity_mapping[a.from_label],'tail':entity_mapping[a.to_label]}
                relations.append(relation)
            data={
                'words':words,
                'entities':entities,
                'relations':relations
            }
        return data


    def export_to_knowledge_graph(self):
        annotations = self.seq_annotations.filter(document_id=self.id)
        relation_annotations = self.seq_relation_annotations.filter(document_id=self.id)
        attribute_annotations = self.seq_attribute_annotations.filter(document_id=self.id)
        words = [token for token in self.text]
        from collections import defaultdict
        entity_result = defaultdict(dict)
        relation_result = defaultdict(set)
        for ent in annotations:
            entity = ''.join(words[ent.start_offset:ent.end_offset])
            if ent.label:
                ent_label = ent.label.text
                if entity not in entity_result[ent_label]:
                    entity_result[ent_label][entity]=dict()
        for rel in relation_annotations:
            to_entity = "".join(words[rel.to_label.start_offset:rel.to_label.end_offset])
            to_entity_type=rel.to_label.label.text
            from_entity = "".join(words[rel.from_label.start_offset:rel.from_label.end_offset])
            from_entity_type=rel.from_label.label.text
            relation = tuple([from_entity_type,from_entity,to_entity_type,to_entity])
            relation_result[rel.relation_label.relation_name].add(relation)
        for attri in attribute_annotations:
            to_label = "".join(words[attri.to_label.start_offset:attri.to_label.end_offset])
            from_label = "".join(words[attri.from_label.start_offset:attri.from_label.end_offset])
            from_label_type=attri.from_label.label.text
            attribute_name = attri.attribute_label.attribute_name
            if attribute_name not in entity_result[from_label_type][from_label].keys():
                entity_result[from_label_type][from_label][attribute_name]=set()
            entity_result[from_label_type][from_label][attribute_name].add(to_label)
        return entity_result, relation_result

    def export_to_json(self):
        entity_result,relation_result=self.export_to_knowledge_graph()
        text=self.text
        return entity_result,relation_result,text

    def make_dataset_json(self):
        return self.make_dataset_for_sequence_labeling_json()

    def make_dataset_for_sequence_labeling_json(self):
        annotations = self.get_annotations()
        entities = [(a.start_offset, a.end_offset, a.label.text) for a in annotations]
        username = annotations[0].user.username
        dataset = {'doc_id': self.id, 'text': self.text, 'entities': entities, 'username': username,
                   'metadata': json.loads(self.metadata)}
        return dataset

    def __str__(self):
        return self.text[:50]


class Annotation(models.Model):
    prob = models.FloatField(default=0.0)
    manual = models.BooleanField(default=False)

    class Meta:
        abstract = True


class AttributeLabel(models.Model):
    attribute_name = models.TextField()
    project=models.ForeignKey(Project,related_name='attribute_labels',on_delete=models.CASCADE)
    label = models.ForeignKey(Label, related_name='attribute_label', on_delete=models.CASCADE)

    def __str__(self):
        # 返回关系名
        return self.attribute_name

    class Meta:
        unique_together = ('attribute_name', 'label','project')
        db_table = 'server_attributeLabel'


class SequenceAnnotation(Annotation):
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    document = models.ForeignKey(Document, related_name='seq_annotations', on_delete=models.CASCADE)
    label = models.ForeignKey(Label, on_delete=models.CASCADE,blank=True,null=True)
    start_offset = models.IntegerField()
    end_offset = models.IntegerField()
    text = models.TextField(default="")

    def clean(self):
        if self.start_offset >= self.end_offset:
            raise ValidationError('start_offset is after end_offset')

    class Meta:
        constraints = [
            UniqueConstraint(fields=['document', 'user', 'start_offset', 'end_offset','label'],
                             name='unique_with_label'),
            UniqueConstraint(fields=['document', 'user', 'start_offset', 'end_offset'],
                             condition=Q(label=None),
                             name='unique_without_label'),
        ]


class RecommendationHistory(models.Model):
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    project = models.ForeignKey(Project, related_name='history', on_delete=models.CASCADE)
    word = models.TextField()
    label = models.ForeignKey(Label, on_delete=models.CASCADE)

    class Meta:
        unique_together = ('user', 'project', 'word', 'label')


class RecommendationRule(models.Model):
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    project = models.ForeignKey(Project, related_name='recommendation_rules', on_delete=models.CASCADE)
    rule = models.TextField()
    label = models.ForeignKey(Label, on_delete=models.CASCADE, blank=True, null=True)

    class Meta:
        constraints = [
            UniqueConstraint(fields=['project', 'user', 'rule', 'label'],
                             name='rule_unique_with_label'),
            UniqueConstraint(fields=['project', 'user', 'rule'],
                             condition=Q(label=None),
                             name='rule_unique_without_label'),
        ]


class TrainedModel(models.Model):
    model = models.CharField(max_length=100, help_text="模型名称", default="")
    dataset = models.CharField(max_length=100, help_text="数据集名称", null=True, blank=True)
    project = models.ForeignKey(Project, on_delete=models.CASCADE, null=True, blank=True)
    model_config = models.TextField(help_text="模型配置信息,为json格式的字符串")
    model_status = models.TextField(help_text="模型状态，训练中/已完成",choices=(
        ("running", "running"), ("finished", "finished"), ("failed", "failed")),
                                  default="running")
    metric = models.TextField(help_text="模型的评价指标结果,为json格式的字符串", default="",blank=True)
    task = models.CharField(max_length=100, default="ner",
                            choices=(("ner", "ner"), ("re", "re"),
                                     ("ner_re", "ner_re")))

    class Meta:
        db_table='server_trainedModel'
        unique_together = (
            ('model', 'dataset','task','project')
        )

    def clean(self):
        # 当 dataset_name 为空时，project 必须存在
        if not self.dataset and not self.project:
            raise ValidationError("未指定 dataset 时，project 是必需的。")


    def save(self, *args, **kwargs):
        self.clean()
        super(TrainedModel, self).save(*args, **kwargs)



class PreSettings(models.Model): # 预标注阶段的设置
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    project = models.ForeignKey(Project, related_name='presettings', on_delete=models.CASCADE)
    preSetPercent = models.FloatField(help_text="预标注的文本的数目")
    preSetAcc = models.FloatField(help_text="预标注阶段需要达到的准确率")
    # def get_setting_serializer(self):
    #     from .serializers import SettingSerializer
    #     return SettingSerializer

    class Meta:
        unique_together = ('user', 'project')

class Setting(models.Model):
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    project = models.ForeignKey(Project, related_name='settings', on_delete=models.CASCADE)
    model_name = models.CharField(max_length=100, help_text="模型名称", default="")
    model_config = models.TextField(help_text="模型训练时的配置信息,为json格式的字符串")
    optimizer_config = models.TextField(help_text="优化器参数配置信息")
    # 推荐方式
    rule=models.BooleanField()
    onlinelearning = models.BooleanField()
    history = models.BooleanField()

    active = models.IntegerField(help_text="主动学习方式")
    acquire = models.IntegerField(help_text="每轮标注规模")

    def get_setting_serializer(self):
        from .serializers import SettingSerializer
        return SettingSerializer

    class Meta:
        unique_together = ('user', 'project')


class AttributeAnnotation(models.Model):
    # id
    # TODO:添加user信息
    document = models.ForeignKey(Document, related_name='seq_attribute_annotations', on_delete=models.CASCADE)
    attribute_label = models.ForeignKey(AttributeLabel, on_delete=models.CASCADE)
    from_label = models.ForeignKey(SequenceAnnotation, related_name='from_attribute', on_delete=models.CASCADE)
    to_label = models.ForeignKey(SequenceAnnotation, related_name='to_attribute', on_delete=models.CASCADE)
    is_recommend = models.BooleanField(default=False)

    def __str__(self):
        # 返回名
        return str(self.from_label) + ',' + str(self.to_label) + "," + str(self.attribute_label)

    class Meta:
        db_table = 'server_attributeAnnotation'
        unique_together = ('document', 'attribute_label', 'from_label', 'to_label')



class RelationLabel(models.Model):
    # id
    relation_name = models.CharField(max_length=100, verbose_name='relation_name')
    head_label=models.ForeignKey(Label,related_name='head_relation_labels',on_delete=models.CASCADE)
    tail_label=models.ForeignKey(Label,related_name='tail_relation_labels',on_delete=models.CASCADE)
    project = models.ForeignKey(Project, related_name='relation_labels', on_delete=models.CASCADE)

    def __str__(self):
        # 返回关系名
        return self.relation_name

    class Meta:
        db_table = 'server_relationLabel'
        unique_together = ('relation_name', 'project','head_label','tail_label')


class RelationAnnotation(models.Model):
    # id
    # TODO:添加user信息
    document = models.ForeignKey(Document, related_name='seq_relation_annotations', on_delete=models.CASCADE)
    relation_label = models.ForeignKey(RelationLabel, on_delete=models.CASCADE)
    from_label = models.ForeignKey(SequenceAnnotation, related_name='from_label', on_delete=models.CASCADE)
    to_label = models.ForeignKey(SequenceAnnotation, related_name='to_label', on_delete=models.CASCADE)
    is_recommend = models.BooleanField(default=False)

    def __str__(self):
        # 返回名
        return str(self.from_label) + ',' + str(self.to_label) + "," + str(self.relation_label)

    class Meta:
        db_table = 'server_relationAnnotation'
        unique_together = ('document', 'relation_label', 'from_label', 'to_label')